package cn.flexui.flexuiserver.controller.web;

import cn.dev33.satoken.stp.StpUtil;
import cn.flexui.flexuiserver.ai.memory.CodeChatMemory;
import cn.flexui.flexuiserver.domain.FlexChatResponse;
import cn.flexui.flexuiserver.domain.constant.AccountConts;
import cn.flexui.flexuiserver.domain.entity.Account;
import cn.flexui.flexuiserver.domain.entity.SessionRecord;
import cn.flexui.flexuiserver.exception.ServiceException;
import cn.flexui.flexuiserver.service.AccountService;
import cn.flexui.flexuiserver.service.ChatRecordService;
import cn.flexui.flexuiserver.service.SessionRecordService;
import cn.flexui.flexuiserver.utils.IdGenerate;
import cn.hutool.core.text.CharSequenceUtil;
import groovy.util.logging.Slf4j;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.math.BigDecimal;

/**
 * @author wanmingyu
 * @create 2024/9/10 20:19
 */
@Slf4j
@Tag(name = "聊天")
@RequestMapping("/chat")
@RestController
public class ChatController {

    @Value("classpath:/prompts/lucide-system-message.st")
    private Resource systemResource;

    private final ChatClient chatClient;
    private final ChatRecordService chatRecordService;
    private final CodeChatMemory codeChatMemory;
    private final AccountService accountService;
    private final SessionRecordService sessionRecordService;

    @Autowired
    public ChatController(ChatRecordService chatRecordService, CodeChatMemory codeChatMemory,
                          AccountService accountService, SessionRecordService sessionRecordService,
                          ChatClient chatClient
    ) {
        this.chatRecordService = chatRecordService;
        this.codeChatMemory = codeChatMemory;
        this.accountService = accountService;
        this.sessionRecordService = sessionRecordService;
        this.chatClient = chatClient;
    }

    @Operation(summary = "生成")
    @Parameter(name = "message", description = "内容")
    @Parameter(name = "fileIds", description = "文件IDS 逗号分割")
    @Parameter(name = "requestId", description = "创建会话时的requestId")
    @GetMapping(value = "/generate", produces = "text/event-stream")
    public Flux<FlexChatResponse> generate(@RequestParam String message,
                                           @RequestParam String requestId,
                                           String fileIds
    ) {

        long userId = StpUtil.getLoginIdAsLong();

        StringBuilder builder = new StringBuilder();
        // 校验用户次数
        Integer accountType = AccountConts.AccountTypeEnum.COUNT_BALANCE.getType();
        Account account = accountService.getByUserIdAndType(userId, accountType);
        if (account.getAccountBalance().compareTo(BigDecimal.ZERO) <= 0) {
            throw new ServiceException(AccountConts.ResponseEnum.BALANCE_INSUFFICIENT);
        }

        SessionRecord sessionRecord = sessionRecordService.getByRequestId(requestId);
        if(sessionRecord == null){
            sessionRecord = new SessionRecord();
            sessionRecord.setRequestId(requestId);
            sessionRecord.setUserId(userId);
            sessionRecord.setSessionName(CharSequenceUtil.sub(message, 0, 10));
            sessionRecordService.save(sessionRecord);
        }

        SimpleLoggerAdvisor loggerAdvisor = new SimpleLoggerAdvisor();
        MessageChatMemoryAdvisor advisor = new MessageChatMemoryAdvisor(codeChatMemory, requestId, 5);
        return chatClient.prompt()
                .advisors(advisor, loggerAdvisor)
                .system(systemResource)
                .user(consumer -> chatRecordService.userChatRecord(consumer, message, fileIds, requestId, userId))
                .stream()
                .chatResponse()
                .map(chatResponse -> FlexChatResponse.mapToFlexChatResponse(chatResponse, builder))
                .doOnComplete(() -> {
                    chatRecordService.assistantChatRecord(requestId, userId, builder);

                    // 消费
                    String outsideCode = IdGenerate.getId();
                    AccountConts.LogDescriptionTypeEnum ldType = AccountConts.LogDescriptionTypeEnum.CONSUME;
                    AccountConts.DirectionTypeEnum typeEnum = AccountConts.DirectionTypeEnum.OUT;
                    accountService.commonUpdateAccount(userId, accountType, BigDecimal.ONE, requestId, outsideCode, ldType, typeEnum, userId + "");
                });
    }

}
